import random
import math

def is_prime(n, k=5):
    """使用Miller-Rabin算法检测素数"""
    if n < 2:
        return False
    if n == 2 or n == 3:
        return True
    if n % 2 == 0:
        return False
    
    # 将n-1写成d * 2^r的形式
    r = 0
    d = n - 1
    while d % 2 == 0:
        r += 1
        d //= 2
    
    # 进行k轮测试
    for _ in range(k):
        a = random.randrange(2, n - 1)
        x = pow(a, d, n)
        
        if x == 1 or x == n - 1:
            continue
        
        for _ in range(r - 1):
            x = pow(x, 2, n)
            if x == n - 1:
                break
        else:
            return False
    
    return True

def generate_prime(bits):
    """生成指定位数的素数"""
    while True:
        # 生成一个奇数
        n = random.getrandbits(bits)
        n |= (1 << bits - 1) | 1  # 设置最高位和最低位为1
        if is_prime(n):
            return n

def gcd(a, b):
    """计算最大公约数"""
    while b:
        a, b = b, a % b
    return a

def extended_gcd(a, b):
    """扩展欧几里得算法"""
    if a == 0:
        return b, 0, 1
    gcd, x1, y1 = extended_gcd(b % a, a)
    x = y1 - (b // a) * x1
    y = x1
    return gcd, x, y

def mod_inverse(e, phi):
    """计算模逆"""
    gcd, x, _ = extended_gcd(e, phi)
    if gcd != 1:
        raise Exception('模逆不存在')
    return (x % phi + phi) % phi

class RSA:
    def __init__(self, key_size=512):
        self.key_size = key_size
        self.generate_keys()
    
    def generate_keys(self):
        """生成密钥对"""
        # 生成两个大素数
        p = generate_prime(self.key_size // 2)
        q = generate_prime(self.key_size // 2)
        
        # 计算n和phi(n)
        self.n = p * q
        phi_n = (p - 1) * (q - 1)
        
        # 选择公钥指数e (通常选择65537)
        self.e = 65537
        if gcd(self.e, phi_n) != 1:
            # 如果e和phi(n)不互质，则选择其他值
            for i in range(self.e, phi_n):
                if gcd(i, phi_n) == 1:
                    self.e = i
                    break
        
        # 计算私钥指数d
        self.d = mod_inverse(self.e, phi_n)
        
        print(f"生成的RSA密钥对:")
        print(f"p = {p}")
        print(f"q = {q}")
        print(f"n = {self.n}")
        print(f"phi(n) = {phi_n}")
        print(f"e = {self.e}")
        print(f"d = {self.d}")
        print()
    
    def encrypt(self, message):
        """加密消息"""
        if isinstance(message, str):
            # 将字符串转换为整数
            message_int = int.from_bytes(message.encode('utf-8'), 'big')
        else:
            message_int = message
        
        if message_int >= self.n:
            raise ValueError("消息太大，无法加密")
        
        encrypted = pow(message_int, self.e, self.n)
        return encrypted
    
    def decrypt(self, encrypted):
        """解密消息"""
        decrypted = pow(encrypted, self.d, self.n)
        
        # 尝试将整数转换回字符串
        try:
            byte_length = (decrypted.bit_length() + 7) // 8
            decrypted_bytes = decrypted.to_bytes(byte_length, 'big')
            return decrypted_bytes.decode('utf-8')
        except:
            return decrypted

def main():
    print("=== RSA算法演示 ===\n")
    
    # 创建RSA实例
    rsa = RSA(key_size=512)
    
    # 原始消息
    original_message = "Hello, RSA Encryption!"
    print(f"原始消息: {original_message}")
    
    # 加密
    encrypted_message = rsa.encrypt(original_message)
    print(f"加密后: {encrypted_message}")
    
    # 解密
    decrypted_message = rsa.decrypt(encrypted_message)
    print(f"解密后: {decrypted_message}")
    
    print(f"\n加密解密成功: {original_message == decrypted_message}")
    
    print("\n=== 数字加密示例 ===")
    # 数字加密示例
    number = 12345
    print(f"原始数字: {number}")
    
    encrypted_number = rsa.encrypt(number)
    print(f"加密后: {encrypted_number}")
    
    decrypted_number = rsa.decrypt(encrypted_number)
    print(f"解密后: {decrypted_number}")
    
    print(f"数字加密解密成功: {number == decrypted_number}")

if __name__ == "__main__":
    main()



